import numpy as np
import os
import scipy.io
import torch
from torch.utils.data import Dataset

class EEGDataset(Dataset):
    def __init__(self, folder_path, window_size=1000, step_size=250, low_freq=1, high_freq=50, fs=250):
        self.folder_path = folder_path
        self.window_size = window_size
        self.step_size = step_size
        self.low_freq = low_freq
        self.high_freq = high_freq
        self.fs = fs
        self.data, self.labels, self.record_ids = self.load_data()

    def bandpass_filter(self, data):
        sos = scipy.signal.butter(5, [self.low_freq, self.high_freq], btype='bandpass', fs=self.fs, output='sos')
        return scipy.signal.sosfilt(sos, data, axis=1)
    def plot_initial_tsne(self):
        # 使用所有数据计算t-SNE
        all_data = np.concatenate(self.data, axis=0)
        labels = np.array(self.labels)
        tsne = TSNE(n_components=2, random_state=42)
        transformed_features = tsne.fit_transform(all_data.reshape(all_data.shape[0], -1))
        plt.figure(figsize=(10, 10))
        scatter = plt.scatter(transformed_features[:, 0], transformed_features[:, 1], c=labels, cmap='viridis', alpha=0.6)
        plt.colorbar(scatter)
        plt.title('Initial t-SNE of EEG Data')
        plt.savefig('initial_tsne.png')
        plt.close()
    def load_data(self):
        data_list, labels_list, record_id_list = [], [], []
        record_id = 0
        for file_name in sorted(os.listdir(self.folder_path)):
            if file_name.endswith('.mat'):
                data_path = os.path.join(self.folder_path, file_name)
                mat = scipy.io.loadmat(data_path)
                data = mat['data']  # Adjust according to your data structure
                label = mat['label']  # Adjust according to your data structure
                for start in range(0, data.shape[1] - self.window_size + 1, self.step_size):
                    end = start + self.window_size
                    segment = data[:, start:end]
                    filtered_segment = self.bandpass_filter(segment)
                    data_list.append(filtered_segment)
                    labels_list.append(label)
                    record_id_list.append(record_id)
                record_id += 1
        return data_list, labels_list, record_id_list

    def __getitem__(self, idx):
        data_segment = self.data[idx]
        label = self.labels[idx]
        data_segment = (data_segment - np.mean(data_segment)) / np.std(data_segment)
        data_segment = np.expand_dims(data_segment, axis=0)
        return torch.tensor(data_segment, dtype=torch.float), torch.tensor(label, dtype=torch.long)

    def __len__(self):
        return len(self.data)

